# 使用说明: deepspeed train_lora.py --deepspeed <$PATH_TO_DEEPSPEED_CONFIG>
#
# 本代码改编自 tatsu-lab@stanford_alpaca，以下是原始版权信息：
#    版权所有 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
#
#    根据 Apache License, Version 2.0 (以下简称“许可证”) 的条款，您不得使用此文件，除非符合许可证的规定。
#    您可以从以下网址获取许可证的副本：
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    除非适用法律要求或书面同意，否则根据许可证分发的软件按“原样”分发，
#    不提供任何明示或暗示的保证或条件。请参阅许可证以了解具体的许可权限和限制。

# 导入必要的库和模块
from dataclasses import dataclass, field
import logging
import pathlib
import typing
import os

from deepspeed import zero
from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
import transformers
from transformers import Trainer, BitsAndBytesConfig, deepspeed
import torch

from fastchat.train.train import (
    DataArguments,
    ModelArguments,
    make_supervised_data_module,
)

from fastchat.train.llama_flash_attn_monkey_patch import (
    replace_llama_attn_with_flash_attn,
)


# 定义训练参数类，继承自 transformers.TrainingArguments
@dataclass
class TrainingArguments(transformers.TrainingArguments):
    """
    扩展训练参数类，添加额外的参数选项。
    参数:
        cache_dir (str): 缓存目录路径。默认为 None。
        optim (str): 优化器类型，默认为 "adamw_torch"。
        model_max_length (int): 模型最大序列长度，默认为 512。
        flash_attn (bool): 是否启用 Flash Attention，默认为 False。
    """
    cache_dir: typing.Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=512,
        metadata={
            "help": "最大序列长度。序列将被右填充（可能被截断）。"
        },
    )
    flash_attn: bool = False


# 定义 LoRA 参数类
@dataclass
class LoraArguments:
    """
    定义 LoRA 参数类。
    参数:
        lora_r (int): LoRA 秩，默认为 8。
        lora_alpha (int): LoRA alpha，默认为 16。
        lora_dropout (float): LoRA dropout 概率，默认为 0.05。
        lora_target_modules (List[str]): 目标模块列表，默认为 ["q_proj", "v_proj"]。
        lora_weight_path (str): LoRA 权重路径，默认为空字符串。
        lora_bias (str): 偏置处理方式，默认为 "none"。
        q_lora (bool): 是否启用 QLoRA，默认为 False。
    """
    lora_r: int = 8
    lora_alpha: int = 16
    lora_dropout: float = 0.05
    lora_target_modules: typing.List[str] = field(
        default_factory=lambda: ["q_proj", "v_proj"]
    )
    lora_weight_path: str = ""
    lora_bias: str = "none"
    q_lora: bool = False


def maybe_zero_3(param):
    """
    处理 DeepSpeed ZeRO-3 的参数状态。
    参数:
        param: 输入参数。
    返回:
        处理后的参数。
    """
    if hasattr(param, "ds_id"):
        assert param.ds_status == ZeroParamStatus.NOT_AVAILABLE
        with zero.GatheredParameters([param]):
            param = param.data.detach().cpu().clone()
    else:
        param = param.detach().cpu().clone()
    return param


# 借鉴自 peft.utils.get_peft_model_state_dict
def get_peft_state_maybe_zero_3(named_params, bias):
    """
    获取 PEFT 模型的状态字典，支持 ZeRO-3。
    参数:
        named_params (dict): 模型参数字典。
        bias (str): 偏置处理方式。
    返回:
        包含 LoRA 参数的状态字典。
    """
    if bias == "none":
        to_return = {k: t for k, t in named_params if "lora_" in k}
    elif bias == "all":
        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
    elif bias == "lora_only":
        to_return = {}
        maybe_lora_bias = {}
        lora_bias_names = set()
        for k, t in named_params:
            if "lora_" in k:
                to_return[k] = t
                bias_name = k.split("lora_")[0] + "bias"
                lora_bias_names.add(bias_name)
            elif "bias" in k:
                maybe_lora_bias[k] = t
        for k, t in maybe_lora_bias:
            if k in lora_bias_names:
                to_return[k] = t
    else:
        raise NotImplementedError
    to_return = {k: maybe_zero_3(v) for k, v in to_return.items()}
    return to_return


def train():
    """
    主训练函数，负责加载模型、数据集、配置并执行训练。
    """
    # 解析命令行参数
    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments, LoraArguments)
    )
    model_args, data_args, training_args, lora_args = parser.parse_args_into_dataclasses()

    # 替换 Flash Attention（如果启用）
    if training_args.flash_attn:
        replace_llama_attn_with_flash_attn()

    # 配置设备映射和分布式训练
    device_map = None
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    ddp = world_size != 1
    if lora_args.q_lora:
        device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)} if ddp else None
        if len(training_args.fsdp) > 0 or deepspeed.is_deepspeed_zero3_enabled():
            logging.warning("FSDP 和 ZeRO3 当前与 QLoRA 不兼容。")

    # 配置计算精度
    compute_dtype = (
        torch.float16
        if training_args.fp16
        else (torch.bfloat16 if training_args.bf16 else torch.float32)
    )

    # 加载预训练模型
    model = transformers.AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        device_map=device_map,
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=compute_dtype,
        ) if lora_args.q_lora else None,
    )

    # 配置 LoRA
    lora_config = LoraConfig(
        r=lora_args.lora_r,
        lora_alpha=lora_args.lora_alpha,
        target_modules=lora_args.lora_target_modules,
        lora_dropout=lora_args.lora_dropout,
        bias=lora_args.lora_bias,
        task_type="CAUSAL_LM",
    )

    # 准备模型进行量化训练
    if lora_args.q_lora:
        model = prepare_model_for_kbit_training(
            model, use_gradient_checkpointing=training_args.gradient_checkpointing
        )
        if not ddp and torch.cuda.device_count() > 1:
            model.is_parallelizable = True
            model.model_parallel = True

    # 应用 LoRA
    model = get_peft_model(model, lora_config)

    # 调整模型中的特定模块到指定精度
    if training_args.flash_attn:
        for name, module in model.named_modules():
            if "norm" in name:
                module = module.to(compute_dtype)
            if "lm_head" in name or "embed_tokens" in name:
                if hasattr(module, "weight"):
                    module = module.to(compute_dtype)

    # 打印可训练参数（仅主进程）
    if training_args.deepspeed is not None and training_args.local_rank == 0:
        model.print_trainable_parameters()

    # 启用梯度检查点（如果需要）
    if training_args.gradient_checkpointing:
        model.enable_input_require_grads()

    # 加载分词器
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="right",
        use_fast=False,
    )
    tokenizer.pad_token = tokenizer.unk_token

    # 创建数据模块
    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)

    # 初始化 Trainer
    trainer = Trainer(
        model=model, tokenizer=tokenizer, args=training_args, **data_module
    )

    # 禁用缓存机制
    model.config.use_cache = False

    # 开始训练（支持从检查点恢复）
    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
        trainer.train(resume_from_checkpoint=True)
    else:
        trainer.train()

    # 保存训练状态
    trainer.save_state()

    # 检查是否启用了 ZeRO-3 模式
    if deepspeed.is_deepspeed_zero3_enabled():
        state_dict_zero3 = trainer.model_wrapped._zero3_consolidated_16bit_state_dict()
        if training_args.local_rank == 0:
            state_dict = state_dict_zero3
    else:
        state_dict = get_peft_state_maybe_zero_3(
            model.named_parameters(), lora_args.lora_bias
        )

    # 保存模型权重
    if training_args.local_rank == 0:
        model.save_pretrained(training_args.output_dir, state_dict=state_dict)


if __name__ == "__main__":
    train()